// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/container/f14_map.h>
#include <cstdint>
#include <mutex>

#include <pollux/common/base/bit_util.h>
#include <pollux/common/base/exceptions.h>

namespace kumo::pollux::cache {
    /// Represents a stream in a table, e.g. nulls/lengths/data of a particular
    /// column. Column-level access tracking uses this to identify the column within
    /// a file or partition. The low 5 bits are the stream kind, e.g. nulls, data
    /// etc. The high 27 bits are the node number in the file schema tree, i.e. the
    /// column.
    class TrackingId {
    public:
        TrackingId() : id_(-1) {
        }

        explicit TrackingId(int32_t id) : id_(id) {
        }

        size_t hash() const {
            return std::hash<int32_t>()(id_);
        }

        bool operator==(const TrackingId &other) const {
            return id_ == other.id_;
        }

        bool empty() const {
            return id_ == -1;
        }

        int32_t id() const {
            return id_;
        }

    private:
        int32_t id_;
    };
} // namespace kumo::pollux::cache

namespace std {
    template<>
    struct hash<::kumo::pollux::cache::TrackingId> {
        size_t operator()(const ::kumo::pollux::cache::TrackingId id) const {
            return id.hash();
        }
    };
} // namespace std

namespace kumo::pollux::cache {
    class FileGroupStats;

    /// Records references and actual uses of a stream.
    struct TrackingData {
        double referencedBytes{};
        double lastReferencedBytes{};
        double readBytes{};
    };

    /// Tracks column access frequency during execution of a query. A ScanTracker is
    /// created at the level of a Task/TableScan, so that all threads of a scan
    /// report in the same tracker. The same ScanTracker tracks all reads of all
    /// partitions of the scan. The groupId argument identifies the file group (e.g.
    /// partition) a tracking event pertains to, since a single ScanTracker can
    /// range over multiple partitions.
    class ScanTracker {
    public:
        ScanTracker() : ScanTracker({}, nullptr, 1) {
        }

        /// Constructs a tracker with 'id'. The tracker will be owned by shared_ptr
        /// and will be referenced from a map from id to weak_ptr to 'this'.
        /// 'unregisterer' is supplied so that the destructor can remove the weak_ptr
        /// from the map of pending trackers. 'loadQuantum' is the largest single IO
        /// size for read.
        ScanTracker(
            std::string_view id,
            std::function<void(ScanTracker *)> unregisterer,
            int32_t loadQuantum,
            FileGroupStats *fileGroupStats = nullptr)
            : id_(id),
              unregisterer_(std::move(unregisterer)),
              fileGroupStats_(fileGroupStats) {
        }

        ~ScanTracker() {
            if (unregisterer_) {
                unregisterer_(this);
            }
        }

        /// Records that a scan references 'bytes' bytes of the stream given by 'id'.
        /// This is called when preparing to read a stripe.
        void recordReference(
            const TrackingId id,
            uint64_t bytes,
            uint64_t fileId,
            uint64_t groupId);

        /// Records that 'bytes' bytes have actually been read from the stream given
        /// by 'id'.
        void recordRead(
            const TrackingId id,
            uint64_t bytes,
            uint64_t fileId,
            uint64_t groupId);

        /// True if 'trackingId' is read at least 'minReadPct' % of the time.
        bool shouldPrefetch(TrackingId id, int32_t minReadPct) {
            return readPct(id) >= minReadPct;
        }

        /// Returns the percentage of referenced columns that are actually read. 100%
        /// if no data.
        int32_t readPct(TrackingId id) {
            std::lock_guard<std::mutex> l(mutex_);
            const auto &data = data_[id];
            if (data.referencedBytes == 0) {
                return 100;
            }
            return data.readBytes / data.referencedBytes * 100;
        }

        TrackingData trackingData(TrackingId id) {
            std::lock_guard<std::mutex> l(mutex_);
            return data_[id];
        }

        std::string_view id() const {
            return id_;
        }

        FileGroupStats *fileGroupStats() const {
            return fileGroupStats_;
        }

        std::string toString() const;

    private:
        // Id of query + scan operator to track.
        const std::string id_;
        const std::function<void(ScanTracker *)> unregisterer_{nullptr};
        FileGroupStats *const fileGroupStats_;

        std::mutex mutex_;
        melon::F14FastMap<TrackingId, TrackingData> data_;
        TrackingData sum_;
    };
} // namespace kumo::pollux::cache
